import jax
import jax.numpy as jnp
from robotisgp import Robotisgp
from learner import *
import defaults

if __name__ == '__main__':
    local_devices_to_use = jax.local_device_count()
    process_count = jax.process_count()
    env_fn = create_fn(Robotisgp)
    core_env = env_fn(
        action_repeat=defaults.action_repeat,
        batch_size=defaults.num_envs // local_devices_to_use // process_count,
        episode_length=defaults.episode_length)
    key = jax.random.PRNGKey(0)
    key, key_rand = jax.random.split(key, 2)
    state = core_env.reset(key_rand)
    step_fn = jax.jit(core_env.step)
    loss = []
    for i in range(1000):
        key, key_act = jax.random.split(key)
        act = jax.random.uniform(key_act,(core_env.batch_size, core_env.action_size))
        state1 = step_fn(state, act)
        p_loss = jnp.mean(jnp.mean(jnp.square(state.obs-state1.obs), axis=-1))
        loss.append(p_loss)
        state = state1
    print(sum(loss)/len(loss))
